import os
import requests
import time
import json
import gradio as gr
from PIL import Image
from io import BytesIO
from datetime import datetime
import threading

# 基础配置
base_url = 'https://api-inference.modelscope.cn/'
api_key = os.environ.get("MODELSCOPE_SDK_TOKEN")
if not api_key:
    raise ValueError("未找到环境变量 MODELSCOPE_SDK_TOKEN，请先设置该变量")

# 保存路径设置
save_dir = os.environ.get("BOZO") + "/Downloads/"
os.makedirs(save_dir, exist_ok=True)  # 确保目录存在

# 通用请求头
common_headers = {
    "Authorization": f"Bearer {api_key}",
    "Content-Type": "application/json",
}

# 尺寸映射
size_mapping = {
    "9:16 (竖屏)": "960x1700",
    "16:9 (横屏)": "1700x960",
    "1:1 (正方形)": "1280x1280",
    "4:3 (传统显示器)": "1480x1100",
    "3:4 (竖向4:3)": "1100x1480",
    "2:3 (超宽屏)": "1080x1536",
    "3:2 (35mm胶片)": "1536x1080",
    "21:9 (超宽屏)": "1980x830",
}

# 模型列表（用于提示）
model_list = [
    "MAILAND/loco-flux-2.5",
    "MAILAND/majicflus_v1",
    "bozoyan/feifei",
    "bozoyan/FK2019_FEI",
    "bozoyan/FK2020_FEI",
    "bozoyan/FK_FEIFEI",
    "bozoyan/F_fei",
    "bozoyan/F_feifei",
    "bozoyan/F_feifei01",
    "bozoyan/FEI_FK",
    "bozoyan/MF_feifei",
    "bozoyan/MEXX_feifei",
    "bozoyan/LOCO_feifei",
    "bozoyan/Xuer_feifei"
]

def generate_image(basetsc, tsc, size_option, model_input):
    # 处理默认值
    if not basetsc.strip():
        basetsc = 'feifei,sexy girl, a photo-realistic portrait of a young asian woman,a young girl,'
    
    if not model_input.strip():
        model_input = "Qwen/Qwen-Image"
    
    # 拼接提示词
    final_prompt = f"{basetsc} {tsc}".strip()
    
    # 获取选中的尺寸
    selected_size = size_mapping[size_option]
    
    try:
        # 提交异步任务
        response = requests.post(
            f"{base_url}v1/images/generations",
            headers={**common_headers, "X-ModelScope-Async-Mode": "true"},
            data=json.dumps({
                "model": model_input,
                "prompt": final_prompt,
                "size": selected_size
            }, ensure_ascii=False).encode('utf-8')
        )
        response.raise_for_status()
        task_id = response.json()["task_id"]
        yield f"任务提交成功，任务ID: {task_id}", None, "生成中..."
        
        # 轮询任务状态
        while True:
            result = requests.get(
                f"{base_url}v1/tasks/{task_id}",
                headers={**common_headers, "X-ModelScope-Task-Type": "image_generation"},
            )
            result.raise_for_status()
            data = result.json()
            
            if data["task_status"] == "SUCCEED":
                # 获取图片URL并处理
                image_url = data["output_images"][0]
                image_response = requests.get(image_url)
                image = Image.open(BytesIO(image_response.content))
                
                # 保存图片到服务器
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"魔搭_{timestamp}.jpg"
                save_path = os.path.join(save_dir, filename)
                image.save(save_path)
                
                yield f"图片生成成功！", image_url, image
                break
                
            elif data["task_status"] == "FAILED":
                yield "图片生成失败！", None, "生成失败"
                break
                
            time.sleep(3)  # 每3秒查询一次状态
            
    except Exception as e:
        yield f"发生错误: {str(e)}", None, "生成失败"

# 创建Gradio界面
with gr.Blocks(title="魔搭图片生成器") as demo:
    gr.Markdown("# 魔搭API图片生成工具")
    
    with gr.Row():
        with gr.Column(scale=1):
            tsc = gr.Textbox(
                label="英文提示词", 
                placeholder="请输入必要的英文提示词",
                lines=20
            )

            basetsc = gr.Textbox(
                label="基础Lora提示词前缀", 
                placeholder="留空使用默认值",
                lines=4
            )
            
            size_option = gr.Dropdown(
                label="图片尺寸",
                choices=list(size_mapping.keys()),
                value="9:16 (竖屏)"
            )
            
            model_input = gr.Textbox(
                label="模型ID", 
                placeholder=f"留空使用默认值 (默认: bozoyan/FEI_FK)\n可用模型: {', '.join(model_list[:3])}...",
                lines=1
            )
            
            generate_btn = gr.Button("生成图片", variant="primary")
            image_url = gr.Textbox(label="图片URL", interactive=False)
            status_text = gr.Textbox(label="状态信息", interactive=False)
        
        with gr.Column(scale=2):
            generated_image = gr.Image(label="生成的图片")
    
    # 设置点击事件
    generate_btn.click(
        fn=generate_image,
        inputs=[basetsc, tsc, size_option, model_input],
        outputs=[status_text, image_url, generated_image]
    )

# 启动服务（在8002端口，允许外部访问）
if __name__ == "__main__":
    demo.launch(server_port=8002, server_name="0.0.0.0")
